module Main where

import Prelude ()
import StdImports
import System.Environment
import System.CPUTime
import System.Timeout
import Text.Printf

import Zeno.Core
import Zeno.HaskellParser
import Zeno.Solver
import Zeno.ProofLogic

filenameDefault = "Z.hs"
conDepthDefault = 1

arg :: String -> IO [String]
arg name = do
  as <- getArgs
  return 
    $ map (\i -> as !! (i + 1))
    $ elemIndices name as 

(<++>) :: (Monoid a, Monad m) => m a -> m a -> m a
x <++> y = (liftM mappend x) `ap` y

main :: IO ()
main = do
  args <- getArgs
  if args == ["!"]
      then test
    else if args == ["comp"]
      then comparison
    else if length args == 1
      then zeno filenameDefault (args !! 0)
    else do
      filepath <- arg "-f"
      [goal] <- arg "-g"
      limit <- arg "-t"
        
      let filepath' = if null filepath 
            then filenameDefault 
            else head filepath

      let computation = zeno filepath' goal
      
      case limit of
        [] -> computation
        [limit_s] -> do
          -- Hard-code 2 secs parsing time
          let limit = ((read limit_s) * (10^6)) + 2 
          maybe_finished <- timeout limit computation
          case maybe_finished of
            Just _ -> return ()
            Nothing ->
              putStrLn $ "Computation timed out (" ++ limit_s ++ "s)."

z :: String -> IO ()
z = zeno filenameDefault

loadZeno :: String -> IO Zeno
loadZeno filename = do
  hsfile <- readFile filename
  create_zeno <- parseHaskellFile filename
  return $ execState create_zeno emptyZeno
  
comparison :: IO () 
comparison = do
  zeno <- loadZeno "Comparison.hs"
  let total = 82
  
  proofs <- forM [1..total] $ \i -> do
    start <- getCPUTime
    
    random <- newStdGen   
    let !proof = runReader (solver ("prop_" ++ show i) random) zeno
    
    end <- getCPUTime
    let diff = (fromIntegral (end - start)) / (10^12)
    
    printf "prop_%i : " (i :: Int)
    if proven proof 
      then printf "%0.3f\n" (diff :: Double)
      else printf "failed\n"
      
    return proof
    
  let count = length $ filter proven proofs
  printf "\nTotal: %i / %i\n" count total
  
test :: IO ()
test = do
  let lemmas = [("prop_dropappend",     2),
                ("prop_zip_append1",    3),
                ("prop_length_remove",  5),
                ("prop_nat_leq_add",    4),
                ("prop_nat_leq_trn",    5),
                ("prop_nat_mul_add",    4),
                ("prop_length_filter",  5),
                ("prop_reverse_filter", 5)]
                
  zeno <- loadZeno filenameDefault
  start <- getCPUTime
  flip mapM_ lemmas $ \(lemma, i) -> do
    random <- newStdGen
    case runReader (solver lemma random) zeno of
      Unknown -> putStrLn $ "Proof failed for " ++ lemma
      Proven {} -> putStrLn $ "Proved " ++ lemma
  end <- getCPUTime
  let diff = (fromIntegral (end - start)) / (10^12)
  printf "\nTotal computation time: %0.3f sec\n" (diff :: Double)
  
meep = print =<< loadZeno filenameDefault
meepcomp = print =<< loadZeno "Comparison.hs" 

zeno :: String -> String -> IO ()
zeno filename lemma = do
  zeno <- loadZeno filename
  -- print zeno

  start <- getCPUTime
  random <- newStdGen
 
  let !proof = runReader (solver lemma random) zeno

  putStrLn $ "\n" ++ showZProof proof
  putStrLn ""
  
  end <- getCPUTime
  
  let diff = (fromIntegral (end - start)) / (10^12)
  printf "\nComputation time: %0.3f sec\n" (diff :: Double)
